"""sync database schema with models

Revision ID: 851830d12fef
Revises: 
Create Date: 2025-07-06 11:30:03.492149

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision: str = '851830d12fef'
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    
    # Drop constraints first from all dependent tables
    op.drop_constraint('conversion_jobs_task_id_fkey', 'conversion_jobs', type_='foreignkey')
    op.drop_constraint('dataflow_quality_metrics_task_id_fkey', 'dataflow_quality_metrics', type_='foreignkey')
    op.drop_constraint('dataflow_results_task_id_fkey', 'dataflow_results', type_='foreignkey')

    # Drop tables that are no longer needed
    op.drop_table('dataflow_quality_metrics')
    op.drop_table('dataflow_results')
    op.drop_table('dataflow_tasks')

    # All operations below are on the 'tasks' table
    op.add_column('tasks', sa.Column('description', sa.Text(), nullable=True, comment='任务描述'))
    op.add_column('tasks', sa.Column('status', sa.Enum('PENDING', 'RUNNING', 'COMPLETED', 'FAILED', 'CANCELLED', name='taskstatus'), nullable=False, server_default='PENDING', comment='任务状态'))
    op.add_column('tasks', sa.Column('library_id', sa.String(length=36), nullable=True, comment='文件库ID'))
    op.add_column('tasks', sa.Column('file_ids', sa.JSON(), nullable=True, comment='处理的文件ID列表'))
    op.add_column('tasks', sa.Column('created_by', sa.String(length=100), nullable=True, comment='创建者'))
    op.add_column('tasks', sa.Column('results', sa.JSON(), nullable=True, comment='处理结果摘要'))
    op.add_column('tasks', sa.Column('quality_metrics', sa.JSON(), nullable=True, comment='质量指标'))
    op.add_column('tasks', sa.Column('current_file', sa.String(length=255), nullable=True, comment='当前处理文件'))
    op.add_column('tasks', sa.Column('total_files', sa.Integer(), nullable=True, comment='总文件数'))
    op.add_column('tasks', sa.Column('processed_files', sa.Integer(), nullable=True, comment='已处理文件数'))
    op.add_column('tasks', sa.Column('failed_files', sa.Integer(), nullable=True, comment='失败文件数'))
    op.add_column('tasks', sa.Column('updated_at', sa.DateTime(), nullable=True, comment='更新时间'))
    op.add_column('tasks', sa.Column('celery_task_id', sa.String(length=255), nullable=True, comment='Celery任务ID'))
    
    op.alter_column('tasks', 'id',
               existing_type=sa.INTEGER(),
               type_=sa.String(length=36),
               existing_nullable=False,
               existing_server_default=sa.text("nextval('tasks_id_seq'::regclass)"))

    op.alter_column('conversion_jobs', 'task_id',
               existing_type=sa.INTEGER(),
               type_=sa.String(length=36),
               existing_nullable=True)

    op.alter_column('tasks', 'name',
               existing_type=sa.VARCHAR(length=255),
               comment='任务名称',
               existing_nullable=False)
    op.alter_column('tasks', 'type',
               existing_type=postgresql.ENUM('PIPELINE_EXECUTION', 'DATA_IMPORT', 'DATA_EXPORT', 'DATA_PROCESSING', 'DOCUMENT_CONVERSION', 'DATASET_GENERATION', name='tasktype'),
               comment='任务类型',
               existing_nullable=False)
    op.alter_column('tasks', 'progress',
               existing_type=sa.INTEGER(),
               comment='进度百分比 0-100',
               existing_nullable=True)
    op.alter_column('tasks', 'config',
               existing_type=postgresql.JSON(astext_type=sa.Text()),
               comment='任务配置',
               existing_nullable=True)
    op.alter_column('tasks', 'result',
               existing_type=postgresql.JSON(astext_type=sa.Text()),
               comment='任务结果',
               existing_nullable=True)
    op.alter_column('tasks', 'error_message',
               existing_type=sa.TEXT(),
               comment='错误信息',
               existing_nullable=True)
    op.alter_column('tasks', 'created_at',
               existing_type=postgresql.TIMESTAMP(),
               comment='创建时间',
               existing_nullable=True)
    op.alter_column('tasks', 'started_at',
               existing_type=postgresql.TIMESTAMP(),
               comment='开始时间',
               existing_nullable=True)
    op.alter_column('tasks', 'completed_at',
               existing_type=postgresql.TIMESTAMP(),
               comment='完成时间',
               existing_nullable=True)
    op.create_foreign_key('fk_tasks_library_id', 'tasks', 'libraries', ['library_id'], ['id'])
    op.create_foreign_key('fk_conversion_jobs_task_id', 'conversion_jobs', 'tasks', ['task_id'], ['id'])
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'tasks', type_='foreignkey')
    op.alter_column('tasks', 'completed_at',
               existing_type=postgresql.TIMESTAMP(),
               comment=None,
               existing_comment='完成时间',
               existing_nullable=True)
    op.alter_column('tasks', 'started_at',
               existing_type=postgresql.TIMESTAMP(),
               comment=None,
               existing_comment='开始时间',
               existing_nullable=True)
    op.alter_column('tasks', 'created_at',
               existing_type=postgresql.TIMESTAMP(),
               comment=None,
               existing_comment='创建时间',
               existing_nullable=True)
    op.alter_column('tasks', 'error_message',
               existing_type=sa.TEXT(),
               comment=None,
               existing_comment='错误信息',
               existing_nullable=True)
    op.alter_column('tasks', 'result',
               existing_type=postgresql.JSON(astext_type=sa.Text()),
               comment=None,
               existing_comment='任务结果',
               existing_nullable=True)
    op.alter_column('tasks', 'config',
               existing_type=postgresql.JSON(astext_type=sa.Text()),
               comment=None,
               existing_comment='任务配置',
               existing_nullable=True)
    op.alter_column('tasks', 'progress',
               existing_type=sa.INTEGER(),
               comment=None,
               existing_comment='进度百分比 0-100',
               existing_nullable=True)
    op.alter_column('tasks', 'type',
               existing_type=postgresql.ENUM('PIPELINE_EXECUTION', 'DATA_IMPORT', 'DATA_EXPORT', 'DATA_PROCESSING', 'DOCUMENT_CONVERSION', 'DATASET_GENERATION', name='tasktype'),
               comment=None,
               existing_comment='任务类型',
               existing_nullable=False)
    op.alter_column('tasks', 'name',
               existing_type=sa.VARCHAR(length=255),
               comment=None,
               existing_comment='任务名称',
               existing_nullable=False)
    op.alter_column('tasks', 'id',
               existing_type=sa.String(length=36),
               type_=sa.INTEGER(),
               existing_nullable=False,
               existing_server_default=sa.text("nextval('tasks_id_seq'::regclass)"))
    op.drop_column('tasks', 'celery_task_id')
    op.drop_column('tasks', 'updated_at')
    op.drop_column('tasks', 'failed_files')
    op.drop_column('tasks', 'processed_files')
    op.drop_column('tasks', 'total_files')
    op.drop_column('tasks', 'current_file')
    op.drop_column('tasks', 'quality_metrics')
    op.drop_column('tasks', 'results')
    op.drop_column('tasks', 'created_by')
    op.drop_column('tasks', 'file_ids')
    op.drop_column('tasks', 'library_id')
    op.drop_column('tasks', 'status')
    op.drop_column('tasks', 'description')
    op.drop_constraint(None, 'dataflow_results', type_='foreignkey')
    op.create_foreign_key('dataflow_results_task_id_fkey', 'dataflow_results', 'dataflow_tasks', ['task_id'], ['id'])
    op.drop_constraint(None, 'dataflow_quality_metrics', type_='foreignkey')
    op.create_foreign_key('dataflow_quality_metrics_task_id_fkey', 'dataflow_quality_metrics', 'dataflow_tasks', ['task_id'], ['id'])

    # Manually added to restore tables in correct order
    op.create_table('dataflow_tasks',
    sa.Column('id', sa.VARCHAR(length=36), autoincrement=False, nullable=False),
    sa.Column('name', sa.VARCHAR(length=255), autoincrement=False, nullable=False, comment='任务名称'),
    sa.Column('description', sa.TEXT(), autoincrement=False, nullable=True, comment='任务描述'),
    sa.Column('pipeline_type', postgresql.ENUM('PRETRAIN_FILTER', 'PRETRAIN_SYNTHETIC', 'SFT_FILTER', 'SFT_SYNTHETIC', name='pipelinetype'), autoincrement=False, nullable=False, comment='流水线类型'),
    sa.Column('library_id', sa.VARCHAR(length=36), autoincrement=False, nullable=True, comment='文件库ID'),
    sa.Column('file_ids', postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True, comment='处理的文件ID列表'),
    sa.Column('created_by', sa.VARCHAR(length=100), autoincrement=False, nullable=True, comment='创建者'),
    sa.Column('config', postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True, comment='流水线配置参数'),
    sa.Column('status', postgresql.ENUM('PENDING', 'RUNNING', 'COMPLETED', 'FAILED', 'CANCELLED', name='taskstatus'), autoincrement=False, nullable=True, comment='任务状态'),
    sa.Column('progress', sa.INTEGER(), autoincrement=False, nullable=True, comment='进度百分比'),
    sa.Column('current_file', sa.VARCHAR(length=255), autoincrement=False, nullable=True, comment='当前处理文件'),
    sa.Column('results', postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True, comment='处理结果摘要'),
    sa.Column('quality_metrics', postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True, comment='质量指标'),
    sa.Column('error_message', sa.TEXT(), autoincrement=False, nullable=True, comment='错误信息'),
    sa.Column('total_files', sa.INTEGER(), autoincrement=False, nullable=True, comment='总文件数'),
    sa.Column('processed_files', sa.INTEGER(), autoincrement=False, nullable=True, comment='已处理文件数'),
    sa.Column('failed_files', sa.INTEGER(), autoincrement=False, nullable=True, comment='失败文件数'),
    sa.Column('created_at', postgresql.TIMESTAMP(), autoincrement=False, nullable=True, comment='创建时间'),
    sa.Column('started_at', postgresql.TIMESTAMP(), autoincrement=False, nullable=True, comment='开始时间'),
    sa.Column('completed_at', postgresql.TIMESTAMP(), autoincrement=False, nullable=True, comment='完成时间'),
    sa.Column('updated_at', postgresql.TIMESTAMP(), autoincrement=False, nullable=True, comment='更新时间'),
    sa.Column('celery_task_id', sa.VARCHAR(length=255), autoincrement=False, nullable=True, comment='Celery任务ID'),
    sa.ForeignKeyConstraint(['library_id'], ['libraries.id'], name='dataflow_tasks_library_id_fkey'),
    sa.PrimaryKeyConstraint('id', name='dataflow_tasks_pkey')
    )
    op.create_table('dataflow_results',
        # ... (Assuming structure from previous context, not fully specified here)
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('task_id', sa.String(36), sa.ForeignKey('dataflow_tasks.id')),
        # ... other columns
    )
    op.create_table('dataflow_quality_metrics',
        # ... (Assuming structure from previous context, not fully specified here)
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('task_id', sa.String(36), sa.ForeignKey('dataflow_tasks.id')),
        # ... other columns
    )
    # ### end Alembic commands ###
